from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
# 修改后的正确导入
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.ollama import Ollama


llm = Ollama(model="deepseek-r1:7b", request_timeout=60.0)

embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")

# 1. 加载文档
documents = SimpleDirectoryReader("data").load_data()

# 2. 使用自定义嵌入模型构建索引（已修复）
index = VectorStoreIndex.from_documents(
    documents, embed_model=embed_model
)

# 3. 创建查询引擎（使用本地LLM）
query_engine = index.as_query_engine(llm=llm)

# 4. 提问
response = query_engine.query("苹果公司的主要业务是什么？")
print("回答:", response)
